{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# DecisionTreeFactor"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/discrete/doc/DecisionTreeFactor.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `DecisionTreeFactor` represents a function over a set of discrete variables, $f(X_1, X_2, ..., X_n)$. It is a versatile building block for representing potentials or probabilities in discrete factor graphs.\n",
        "\n",
        "Internally, it uses an `AlgebraicDecisionTree` (ADT) to store the function's values. This representation can be very efficient, especially for sparse factors where many assignments have the same value (e.g., zero).\n",
        "\n",
        "`DecisionTreeFactor` is the base class for more specialized factors like `DiscreteConditional`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "import graphviz\n",
        "\n",
        "from gtsam.symbol_shorthand import A, B"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a DecisionTreeFactor\n",
        "\n",
        "A factor is defined by its `DiscreteKeys` (the variables it depends on) and a table of values. The table specifies the factor's output for every possible assignment of its variables. The values in the table are ordered such that the last key varies fastest."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Factor f1 from string ---\n",
            "DecisionTreeFactor\n",
            " f[ (a0,2), (b0,2), ]\n",
            " Choice(b0) \n",
            " 0 Choice(a0) \n",
            " 0 0 Leaf    1\n",
            " 0 1 Leaf    3\n",
            " 1 Choice(a0) \n",
            " 1 0 Leaf    2\n",
            " 1 1 Leaf    4\n",
            "\n",
            "--- Factor f2 from vector ---\n",
            "DecisionTreeFactor\n",
            " f[ (a0,2), ]\n",
            " Choice(a0) \n",
            " 0 Leaf  0.8\n",
            " 1 Leaf  0.2\n"
          ]
        }
      ],
      "source": [
        "# Define keys for two binary variables\n",
        "KeyA = (A(0), 2)\n",
        "KeyB = (B(0), 2)\n",
        "\n",
        "# --- Method 1: From a spec string ---\n",
        "# The values correspond to assignments (A,B) in the order:\n",
        "# (0,0), (0,1), (1,0), (1,1)\n",
        "f1_string = gtsam.DecisionTreeFactor([KeyA, KeyB], \"1 2 3 4\")\n",
        "print(\"--- Factor f1 from string ---\")\n",
        "f1_string.print()\n",
        "\n",
        "# --- Method 2: From a list of values ---\n",
        "f2_vector = gtsam.DecisionTreeFactor([KeyA], [0.8, 0.2])\n",
        "print(\"\\n--- Factor f2 from vector ---\")\n",
        "f2_vector.print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operations on DecisionTreeFactor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Value of f1 at A=1, B=0: 3.0\n",
            "\n",
            "--- Product Factor f3 = f1 * f2 ---\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<table class='DecisionTreeFactor'>\n",
              "  <thead>\n",
              "    <tr><th>a0</th><th>b0</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><th>0</th><td>0.8</td></tr>\n",
              "    <tr><th>0</th><th>1</th><td>1.6</td></tr>\n",
              "    <tr><th>1</th><th>0</th><td>0.6</td></tr>\n",
              "    <tr><th>1</th><th>1</th><td>0.8</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/markdown": [
              "|a0|b0|value|\n",
              "|:-:|:-:|:-:|\n",
              "|0|0|0.8|\n",
              "|0|1|1.6|\n",
              "|1|0|0.6|\n",
              "|1|1|0.8|\n"
            ],
            "text/plain": [
              "DecisionTreeFactor\n",
              " f[ (a0,2), (b0,2), ]\n",
              " Choice(b0) \n",
              " 0 Choice(a0) \n",
              " 0 0 Leaf  0.8\n",
              " 0 1 Leaf  0.6\n",
              " 1 Choice(a0) \n",
              " 1 0 Leaf  1.6\n",
              " 1 1 Leaf  0.8"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# --- Evaluate ---\n",
        "# Get the factor's value for a specific assignment.\n",
        "assignment = gtsam.DiscreteValues()\n",
        "assignment[A(0)] = 1\n",
        "assignment[B(0)] = 0\n",
        "\n",
        "value = f1_string(assignment)\n",
        "print(f\"Value of f1 at A=1, B=0: {value}\")\n",
        "\n",
        "# --- Multiplication ---\n",
        "# Multiplying factors corresponds to the product of their functions.\n",
        "# The resulting factor is over the union of their variables.\n",
        "# f3(A,B) = f1(A,B) * f2(A)\n",
        "f3_product = f1_string * f2_vector\n",
        "print(\"\\n--- Product Factor f3 = f1 * f2 ---\")\n",
        "f3_product"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualization\n",
        "\n",
        "We can visualize the underlying Algebraic Decision Tree structure. This shows how the factor's values are stored efficiently. Each path from the root to a leaf corresponds to an assignment, and the leaf holds the value."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Title: G Pages: 1 -->\n",
              "<svg width=\"278pt\" height=\"185pt\"\n",
              " viewBox=\"0.00 0.00 278.00 184.70\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 180.7)\">\n",
              "<title>G</title>\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-180.7 274,-180.7 274,4 -4,4\"/>\n",
              "<!-- 0x135734e58 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>0x135734e58</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"135\" cy=\"-156.57\" rx=\"20.13\" ry=\"20.13\"/>\n",
              "<text text-anchor=\"middle\" x=\"135\" y=\"-151.52\" font-family=\"Times,serif\" font-size=\"14.00\">b0</text>\n",
              "</g>\n",
              "<!-- 0x135734f08 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>0x135734f08</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-80.82\" rx=\"19.62\" ry=\"19.62\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-75.77\" font-family=\"Times,serif\" font-size=\"14.00\">a0</text>\n",
              "</g>\n",
              "<!-- 0x135734e58&#45;&gt;0x135734f08 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>0x135734e58&#45;&gt;0x135734f08</title>\n",
              "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M126.47,-138.09C122.22,-129.4 117,-118.69 112.25,-108.95\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"115.54,-107.73 108.01,-100.28 109.25,-110.8 115.54,-107.73\"/>\n",
              "</g>\n",
              "<!-- 0x1357353a8 -->\n",
              "<g id=\"node5\" class=\"node\">\n",
              "<title>0x1357353a8</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"171\" cy=\"-80.82\" rx=\"19.62\" ry=\"19.62\"/>\n",
              "<text text-anchor=\"middle\" x=\"171\" y=\"-75.77\" font-family=\"Times,serif\" font-size=\"14.00\">a0</text>\n",
              "</g>\n",
              "<!-- 0x135734e58&#45;&gt;0x1357353a8 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>0x135734e58&#45;&gt;0x1357353a8</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M143.53,-138.09C147.78,-129.4 153,-118.69 157.75,-108.95\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"160.75,-110.8 161.99,-100.28 154.46,-107.73 160.75,-110.8\"/>\n",
              "</g>\n",
              "<!-- 0x147a5f210 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>0x147a5f210</title>\n",
              "<polygon fill=\"none\" stroke=\"black\" points=\"54,-25.2 0,-25.2 0,0 54,0 54,-25.2\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-7.55\" font-family=\"Times,serif\" font-size=\"14.00\">1.00</text>\n",
              "</g>\n",
              "<!-- 0x135734f08&#45;&gt;0x147a5f210 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>0x135734f08&#45;&gt;0x147a5f210</title>\n",
              "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M84.78,-66.74C74.32,-57.13 59.95,-43.91 48.16,-33.06\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"50.59,-30.54 40.86,-26.35 45.85,-35.69 50.59,-30.54\"/>\n",
              "</g>\n",
              "<!-- 0x135734f60 -->\n",
              "<g id=\"node4\" class=\"node\">\n",
              "<title>0x135734f60</title>\n",
              "<polygon fill=\"none\" stroke=\"black\" points=\"126,-25.2 72,-25.2 72,0 126,0 126,-25.2\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-7.55\" font-family=\"Times,serif\" font-size=\"14.00\">3.00</text>\n",
              "</g>\n",
              "<!-- 0x135734f08&#45;&gt;0x135734f60 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>0x135734f08&#45;&gt;0x135734f60</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M99,-61.02C99,-53.49 99,-44.74 99,-36.82\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"102.5,-36.87 99,-26.87 95.5,-36.87 102.5,-36.87\"/>\n",
              "</g>\n",
              "<!-- 0x135734520 -->\n",
              "<g id=\"node6\" class=\"node\">\n",
              "<title>0x135734520</title>\n",
              "<polygon fill=\"none\" stroke=\"black\" points=\"198,-25.2 144,-25.2 144,0 198,0 198,-25.2\"/>\n",
              "<text text-anchor=\"middle\" x=\"171\" y=\"-7.55\" font-family=\"Times,serif\" font-size=\"14.00\">2.00</text>\n",
              "</g>\n",
              "<!-- 0x1357353a8&#45;&gt;0x135734520 -->\n",
              "<g id=\"edge5\" class=\"edge\">\n",
              "<title>0x1357353a8&#45;&gt;0x135734520</title>\n",
              "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M171,-61.02C171,-53.49 171,-44.74 171,-36.82\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"174.5,-36.87 171,-26.87 167.5,-36.87 174.5,-36.87\"/>\n",
              "</g>\n",
              "<!-- 0x135734f90 -->\n",
              "<g id=\"node7\" class=\"node\">\n",
              "<title>0x135734f90</title>\n",
              "<polygon fill=\"none\" stroke=\"black\" points=\"270,-25.2 216,-25.2 216,0 270,0 270,-25.2\"/>\n",
              "<text text-anchor=\"middle\" x=\"243\" y=\"-7.55\" font-family=\"Times,serif\" font-size=\"14.00\">4.00</text>\n",
              "</g>\n",
              "<!-- 0x1357353a8&#45;&gt;0x135734f90 -->\n",
              "<g id=\"edge6\" class=\"edge\">\n",
              "<title>0x1357353a8&#45;&gt;0x135734f90</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M185.22,-66.74C195.68,-57.13 210.05,-43.91 221.84,-33.06\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"224.15,-35.69 229.14,-26.35 219.41,-30.54 224.15,-35.69\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x136ccbc80>"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "graphviz.Source(f1_string.dot())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Rich Display in Jupyter\n",
        "In a notebook, factors are displayed as easy-to-read tables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<table class='DecisionTreeFactor'>\n",
              "  <thead>\n",
              "    <tr><th>a0</th><th>b0</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><th>0</th><td>1</td></tr>\n",
              "    <tr><th>0</th><th>1</th><td>2</td></tr>\n",
              "    <tr><th>1</th><th>0</th><td>3</td></tr>\n",
              "    <tr><th>1</th><th>1</th><td>4</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/markdown": [
              "|a0|b0|value|\n",
              "|:-:|:-:|:-:|\n",
              "|0|0|1|\n",
              "|0|1|2|\n",
              "|1|0|3|\n",
              "|1|1|4|\n"
            ],
            "text/plain": [
              "DecisionTreeFactor\n",
              " f[ (a0,2), (b0,2), ]\n",
              " Choice(b0) \n",
              " 0 Choice(a0) \n",
              " 0 0 Leaf    1\n",
              " 0 1 Leaf    3\n",
              " 1 Choice(a0) \n",
              " 1 0 Leaf    2\n",
              " 1 1 Leaf    4"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "f1_string"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}